# coding: utf-8

import contextlib
import functools
import threading

import pymysql
from DBUtils.PersistentDB import PersistentDB

from .config import PooledDBConfig


class Transaction:
    """事务类"""

    def __init__(self, conn):
        self.conn = conn
        self.local_stack = threading.local()

    @property
    def context(self):
        if hasattr(self.local_stack, 'transaction'):
            return self.local_stack.transaction
        self.local_stack.transaction = False
        return self.local_stack.transaction

    @context.setter
    def context(self, x):
        self.local_stack.transaction = x

    def begin(self, is_select=False):
        if is_select:
            return

        if not self.context:
            self.conn.begin()

        self.context += 1

    def commit(self, is_select=False):
        if is_select:
            return

        if self.context:
            self.context -= 1

        if not self.context:
            self.conn.commit()

    def rollback(self, is_select=False):
        if is_select:
            return
        if self.context:
            self.context -= 1

        if not self.context:
            self.conn.rollback()

    def clear(self):
        if self.context:
            self.conn.commit()
        self.context = False


class SQLBase:
    """SQL Base class, init Database connection"""

    def __init__(self, host, database, port, user, passwd, **kw):
        _params = PooledDBConfig.__dict__.copy()
        for k in PooledDBConfig.__dict__:
            if k.startswith('__') or k.startswith('_'):
                _params.pop(k)
        self.host = _params['host'] = host
        self.db = _params['db'] = database
        self.port = _params['port'] = int(port)
        self.user = _params['user'] = user
        self.passwd = _params['passwd'] = passwd
        _params.update(kw)
        self._db_config = _params
        self._pool = PersistentDB(pymysql, **_params)
        self._conn = self._pool.connection()
        self._cursor = self._conn.cursor()
        self._transaction = Transaction(self._conn)

    @property
    def connection(self):
        if not all([self._conn, self._cursor]):
            self._conn = self._pool.connection()
            self._cursor = self._conn.cursor()
        return self._conn

    @property
    def transaction(self):
        if not self._transaction:
            self._transaction = Transaction(self.connection)
        return self._transaction

    @property
    def cursor(self):
        if not self._cursor:
            self.connection
        return self._cursor

    def begin(self):
        self.transaction.begin()

    def commit(self):
        self.transaction.commit()

    def rollback(self):
        self.transaction.rollback()

    def close(self):
        self.cursor.close()
        self.conn.close()

        self._conn = None
        self._cursor = None
        self._transaction = None

    def execute(self, sql, values: tuple=None, select=False):
        try:
            if not select:
                self.begin()
            if values:
                rows = self.cursor.execute(sql, values)
            else:
                rows = self.cursor.execute(sql)
            if not select:
                self.commit()
        except:
            if not select:
                self.rollback()
            raise
        return rows

    def with_transaction(self, func):
        @functools.wraps(func)
        def deco(*args, **kw):
            try:
                self.begin()
                ret = func(*args, **kw)
                self.commit()
            except Exception as e:
                self.rollback()
                raise
            return ret
        return deco

    @contextlib.contextmanager
    def transaction_context(self):
        try:
            self.begin()
            yield
            self.commit()
        except:
            self.rollback()
            raise

    def yield_data(self, block=1000):
        # return self.cursor.fetchall()
        while True:
            ret = self.cursor.fetchmany(block)
            if not ret:
                break
            for r in ret:
                yield r
